package com.element.oauth2.server.global.filter;

import com.element.oauth2.constant.SecurityParams;
import com.element.oauth2.server.handler.AuthErrorAuthenticationFailureHandler;
import org.springframework.http.HttpMethod;
import org.springframework.security.authentication.AbstractAuthenticationToken;
import org.springframework.security.authentication.AuthenticationManager;
import org.springframework.security.authentication.UsernamePasswordAuthenticationToken;
import org.springframework.security.core.Authentication;
import org.springframework.security.oauth2.core.OAuth2AuthenticationException;
import org.springframework.security.oauth2.server.authorization.authentication.OAuth2AuthorizationCodeRequestAuthenticationToken;
import org.springframework.security.oauth2.server.authorization.authentication.OAuth2AuthorizationConsentAuthenticationToken;
import org.springframework.security.web.DefaultRedirectStrategy;
import org.springframework.security.web.authentication.AuthenticationConverter;
import org.springframework.security.web.authentication.AuthenticationSuccessHandler;
import org.springframework.security.web.authentication.WebAuthenticationDetailsSource;
import org.springframework.security.web.util.RedirectUrlBuilder;
import org.springframework.security.web.util.UrlUtils;
import org.springframework.security.web.util.matcher.AntPathRequestMatcher;
import org.springframework.security.web.util.matcher.OrRequestMatcher;
import org.springframework.security.web.util.matcher.RequestMatcher;
import org.springframework.util.StringUtils;
import org.springframework.web.filter.OncePerRequestFilter;
import org.springframework.web.util.UriComponentsBuilder;
import org.springframework.web.util.UriUtils;
import reactor.util.annotation.NonNull;

import javax.servlet.FilterChain;
import javax.servlet.ServletException;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;
import java.io.IOException;
import java.nio.charset.StandardCharsets;
import java.util.Objects;
import java.util.Set;

public class OAuthAuthorizationEndpointFilter extends OncePerRequestFilter {

    private final RequestMatcher authorizationEndpointMatcher = createDefaultRequestMatcher();

    private final AuthenticationSuccessHandler authenticationSuccessHandler = this::sendAuthorizationResponse;

    private final AuthenticationManager authenticationManager;

    private final AuthenticationConverter authenticationConverter;

    public OAuthAuthorizationEndpointFilter(AuthenticationManager authenticationManager,
                                            AuthenticationConverter authenticationConverter) {
        this.authenticationManager = authenticationManager;
        this.authenticationConverter = authenticationConverter;
    }

    @Override
    protected void doFilterInternal(@NonNull HttpServletRequest request,
                                    @NonNull HttpServletResponse response,
                                    @NonNull FilterChain filterChain) throws ServletException, IOException {
        if (!this.authorizationEndpointMatcher.matches(request)) {
            filterChain.doFilter(request, response);
        } else {
            try {
                Authentication authentication = this.authenticationConverter.convert(request);
                if (authentication instanceof AbstractAuthenticationToken) {
                    ((AbstractAuthenticationToken) authentication).setDetails(new WebAuthenticationDetailsSource().buildDetails(request));
                }
                Authentication authenticationResult = authenticationManager.authenticate(authentication);
                if (null == authenticationResult || !authenticationResult.isAuthenticated()) {
                    filterChain.doFilter(request, response);
                    return;
                }
                if (authenticationResult instanceof OAuth2AuthorizationConsentAuthenticationToken) {
                    assert authentication instanceof OAuth2AuthorizationCodeRequestAuthenticationToken;
                    this.sendAuthorizationConsent(request, response, (OAuth2AuthorizationCodeRequestAuthenticationToken) authentication, (OAuth2AuthorizationConsentAuthenticationToken) authenticationResult);
                } else if (authenticationResult instanceof UsernamePasswordAuthenticationToken) {
                    filterChain.doFilter(request, response);
                } else {
                    authenticationSuccessHandler.onAuthenticationSuccess(request, response, authenticationResult);
                }
            } catch (OAuth2AuthenticationException var6) {
                new AuthErrorAuthenticationFailureHandler().onAuthenticationFailure(request, response, var6);
            }

        }
    }

    private RequestMatcher createDefaultRequestMatcher() {
        RequestMatcher authorizationRequestGetMatcher = new AntPathRequestMatcher(SecurityParams.AUTHORIZATION_ENDPOINT_URL, HttpMethod.GET.name());
        RequestMatcher authorizationRequestPostMatcher = new AntPathRequestMatcher(SecurityParams.AUTHORIZATION_ENDPOINT_URL, HttpMethod.POST.name());
        RequestMatcher authorizationRequestMatcher = new OrRequestMatcher(authorizationRequestGetMatcher, authorizationRequestPostMatcher);
        return new OrRequestMatcher(authorizationRequestMatcher);
    }

    private void sendAuthorizationConsent(HttpServletRequest request,
                                          HttpServletResponse response,
                                          OAuth2AuthorizationCodeRequestAuthenticationToken authorizationCodeRequestAuthentication,
                                          OAuth2AuthorizationConsentAuthenticationToken authorizationConsentAuthentication) throws IOException {
        String clientId = authorizationConsentAuthentication.getClientId();
        Set<String> requestedScopes = authorizationCodeRequestAuthentication.getScopes();
        String state = authorizationConsentAuthentication.getState();
        String redirectUri = UriComponentsBuilder.fromUriString(this.resolveConsentUri(request)).queryParam(SecurityParams.SCOPE, String.join(" ", requestedScopes)).queryParam(SecurityParams.CLIENT_ID, clientId).queryParam(SecurityParams.STATE, state).toUriString();
        new DefaultRedirectStrategy().sendRedirect(request, response, redirectUri);
    }

    private void sendAuthorizationResponse(HttpServletRequest request,
                                           HttpServletResponse response,
                                           Authentication authentication) throws IOException {
        OAuth2AuthorizationCodeRequestAuthenticationToken authorizationCodeRequestAuthentication = (OAuth2AuthorizationCodeRequestAuthenticationToken) authentication;
        UriComponentsBuilder uriBuilder = UriComponentsBuilder.fromUriString(Objects.requireNonNull(authorizationCodeRequestAuthentication.getRedirectUri())).queryParam("code", Objects.requireNonNull(authorizationCodeRequestAuthentication.getAuthorizationCode()).getTokenValue());
        if (StringUtils.hasText(authorizationCodeRequestAuthentication.getState())) {
            uriBuilder.queryParam(SecurityParams.STATE, UriUtils.encode(authorizationCodeRequestAuthentication.getState(), StandardCharsets.UTF_8));
        }

        String redirectUri = uriBuilder.build(true).toUriString();
        new DefaultRedirectStrategy().sendRedirect(request, response, redirectUri);
    }

    private String resolveConsentUri(HttpServletRequest request) {
        if (UrlUtils.isAbsoluteUrl(SecurityParams.CONSENT_PAGE)) {
            return SecurityParams.CONSENT_PAGE;
        } else {
            RedirectUrlBuilder urlBuilder = new RedirectUrlBuilder();
            urlBuilder.setScheme(request.getScheme());
            urlBuilder.setServerName(request.getServerName());
            urlBuilder.setPort(request.getServerPort());
            urlBuilder.setContextPath(request.getContextPath());
            urlBuilder.setPathInfo(SecurityParams.CONSENT_PAGE);
            return urlBuilder.getUrl();
        }
    }
}